'''
Module to download Equity Bhavcopy from NSE
Created on Jul 12, 2014
@author: ANAND
'''
import os
import re
import zipfile
import urllib2

def unzip(source_filename, dest_dir):
    """Function to unzip downloaded bhavcopy zip file.
    Input source file name (full path) and destination path
    Returns list of file names extracted from zip.
    """
    with zipfile.ZipFile(source_filename) as zf:
        for member in zf.infolist():
            words = member.filename.split('/')
            path = dest_dir
            for word in words[:-1]:
                #drive, word = os.path.splitdrive(word)
                #head, word = os.path.split(word)
                word = os.path.splitdrive(word)[1]
                word = os.path.split(word)[1]
                if word in (os.curdir, os.pardir, ''): continue
                path = os.path.join(path, word)
            zf.extract(member, path)
    return(words)

def bhav_downloader(url, dest, my_date):
    """Function for Bhavcopy download from NSEIndia site.
    Uses urllib2 to dowanload from http link.
    Takes input URL of file to download, destination to save and date of file data.
    Calls unzip function.
    """
    res = 0
    file_name = url.split('/')[-1]
    newfile = my_date + "_eb.dat"
    
    hdr = {'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.11 (KHTML, like Gecko) Chrome/23.0.1271.64 Safari/537.11',
           'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
           'Accept-Charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.3',
           'Accept-Encoding': 'none',
           'Accept-Language': 'en-US,en;q=0.8',
           'Connection': 'keep-alive'}
    
    req = urllib2.Request(url, headers=hdr)
    
    try:
        bhav1 = urllib2.urlopen(req)
    except urllib2.HTTPError, e:
        print e.fp.read()
    
    try:
        meta = bhav1.info()
    except UnboundLocalError:
        print "File does not exist.\nPlease Enter correct date."
        return(res, newfile)
    
    file_size = int(meta.getheaders("Content-Length")[0])
    print "Downloading: %s Size: %s Bytes" % (file_name, file_size)
    
    data = bhav1.read()
    file_size_dl = len(data)
    
    file_name = dest + "/" + file_name
    
    out_bhav1 = open(file_name, "wb")
    out_bhav1.write(data)
    out_bhav1.close()
    
    if file_size_dl == file_size:
        print "Download complete!"
    
    bhav_file = unzip(file_name, dest)
    os.remove(file_name)
    oldfile = bhav_file[0]
    
    try: os.remove(newfile)
    except: pass
    os.rename(oldfile, newfile)
    res = 1
    return(res, newfile)

def my_downloader(cur_date, pre_date, dest):
    """"
    Takes two dates, for which files to download and destination location to save.
    Customize URL accordingly.
    Calls bhav_download function.
    """
    # my_date = 01-07-2014
    # my_months = {"01":"JAN", "02":"FEB", "03":"MAR", "04":"APR", "05":"MAY", "06":"JUN", "07":"JUL", "08":"AUG", "09":"SEP", "10":"OCT", "11":"NOV", "12":"DEC"}
    # my_url = "http://www.nseindia.com/content/historical/EQUITIES/<YYYY>/<MON>/cm<DD><MON><YYYY>bhav.csv.zip"
    #
    # e.g. my_url = "http://www.nseindia.com/content/historical/EQUITIES/2014/JUL/cm10JUL2014bhav.csv.zip"
    res = 0
    two_dates = (cur_date, pre_date)
    my_months = {"01":"JAN", "02":"FEB", "03":"MAR", "04":"APR", "05":"MAY", "06":"JUN", "07":"JUL", "08":"AUG", "09":"SEP", "10":"OCT", "11":"NOV", "12":"DEC"}
    bhav_files = []
    os.chdir(dest)
    for my_date in two_dates:
        my_dd, my_mm, my_yyyy = my_date.split("-")
        my_mon = my_months[my_mm]
        
        my_url = "http://www.nseindia.com/content/historical/EQUITIES/<YYYY>/<MON>/cm<DD><MON><YYYY>bhav.csv.zip"
        my_url = re.sub("<YYYY>",my_yyyy, my_url)
        my_url = re.sub("<MON>",my_mon, my_url)
        my_url = re.sub("<DD>",my_dd, my_url)
        
        bhav_file = bhav_downloader(my_url, dest, my_date)
        if bhav_file[0] == 1:
            bhav_files.append(bhav_file[1])
        else:
            return(res, bhav_files)
    res = 1
    return(res, bhav_files)

